# 这里需要一个方法, 参数是一个pytorch模型, 我需要输出这个模型的各种各样的指标

import torch
from torch import dtype
from torch.profiler import profile, record_function, ProfilerActivity
import torch.cuda
import os

from torchsummary import summary
from transformers import AutoConfig, AutoTokenizer

from base_model.nlp.base_model.first_model_mla import UsherConfig, Transformer

# 注册自定义配置类
AutoConfig.register("usherTransformer", UsherConfig)
torch.backends.cuda.matmul.allow_tf32 = True  # 启用TF32加速
torch.set_float32_matmul_precision('medium')  # 优化矩阵乘法
# 这个设置用于调试CUDA操作。当设置为1时，CUDA操作会以同步模式运行，这意味着每个CUDA操作在返回之前都会等待完成。这有助于更容易地捕获和定位CUDA相关错误。
# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
# 这个设置配置了PyTorch的CUDA内存分配策略。expandable_segments:True允许CUDA内存段扩展，而不是预先分配固定大小的内存块。这种设置可以提高内存使用效率，尤其是在处理动态大小的数据集或模型时。
# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
# 这个设置启用了CUDA的直接共享访问（Direct Shared Access, DSA）。DSA允许不同的CUDA上下文之间更高效地共享内存资源，从而可能提升多GPU环境下的性能。
os.environ['TORCH_USE_CUDA_DSA'] = '1'


def print_model_stats(model, input_sample=None):
    """输出模型统计信息"""

    # 1. 参数统计
    def print_params():
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"总参数量：{total_params:,}")
        print(f"可训练参数量：{trainable_params:,}")

    # 2. 模型大小（MB）
    def print_size():
        param_size = sum(p.nelement() * p.element_size() for p in model.parameters())
        buffer_size = sum(b.nelement() * b.element_size() for b in model.buffers())
        total_size = (param_size + buffer_size) / 1024 ** 2
        print(f"模型大小：{total_size:.2f}MB")

    # 3. FLOPs统计
    def print_flops():
        try:
            from thop import profile as thop_profile
            macs, params = thop_profile(model, inputs=(input_sample,))
            print(f"FLOPs：{macs / 1e9:.2f}G")
        except ImportError:
            print("请安装thop库以计算FLOPs")

    # 4. 内存使用统计
    def print_memory():
        if torch.cuda.is_available():
            print("CUDA内存使用情况：")
            print(f"已分配：{torch.cuda.memory_allocated() / 1e6:.2f}MB")
            print(f"已保留：{torch.cuda.memory_reserved() / 1e6:.2f}MB")
        else:
            print("CPU内存统计：不支持")

    # 5. 层级结构
    def print_layers(input_sample):
        print("模型层级结构：")
        summary(model, input_data=input_sample)

    # 6. 计算图
    def plot_graph():
        try:
            from torchviz import make_dot
            y = model(input_sample)
            dot = make_dot(y)
            abspath = os.path.abspath("result/model_graph.gv")
            # 如果没有文件，则创建一个空的文件
            if not os.path.exists(abspath):
                with open(abspath, 'w') as f:
                    f.write("")
                # 为关键节点添加注释（示例）
            for node in dot.body:
                if "attention.matmul" in node:
                    node.label = "Attention Matrix Multiply"
                elif "linear.bias" in node:
                    node.label = "Linear Layer Bias Addition"
            dot.render(abspath, view=False)
            print("计算图已保存为 model_graph.gv.pdf")
        except ImportError:
            print("请安装torchviz和graphviz库以生成计算图")

    # 7. 性能分析
    def profile_performance():
        with profile(activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU], profile_memory=True, record_shapes=True) as prof:
            with record_function("模型推理"):
                model(input_sample)
        if torch.cuda.is_available():
            print("性能分析结果：")
            print(prof.key_averages().table(sort_by="cuda_time_total"))
        else:
            print("CPU性能分析：未实现")

    # 执行所有统计
    print("模型统计报告：")
    print_params()
    print_size()
    if input_sample is not None:
        print_flops()
        profile_performance()
        plot_graph()
    print_memory()
    print_layers(input_sample)


# 加载分词器
config = UsherConfig()
# 配置
config = AutoConfig.from_pretrained(config.path)
print(config)

# 分词器
tokenizer = AutoTokenizer.from_pretrained(config.path)


# 数据预处理函数
def preprocess_function(examples):
    inputs = examples["input"]
    outputs = examples["output"]

    prompts = [
        f"<｜begin_of_sentence｜><｜User｜>{input}<｜Assistant｜>{output}<｜end｜>"
        for input, output in zip(inputs, outputs)
    ]

    tokenized = tokenizer(prompts, max_length=config.max_seq_len, truncation=True, padding="max_length", return_tensors="pt")

    # 生成 labels：input_ids 右偏移一位，最后一个位置设为 -100
    labels = tokenized["input_ids"][:, 1:].clone().detach()  # 截取从第二个 token 开始
    padding = torch.full((labels.shape[0], 1), tokenizer.pad_token_id, dtype=torch.long)  # 补充 填充值 到末尾
    tokenized["check_labels"] = torch.cat([labels, padding], dim=1)

    return tokenized


if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 使用示例：
    # 初始化模型并移动到指定设备
    model = Transformer(config).to(device)
    # 加载权重，并确保权重加载到正确的设备上
    model_path = config.path + r"\model.pth"
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found at {model_path}")

    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)

    input_sample = preprocess_function({"input": ["2749045"], "output": ["二七四九零四五"]})

    print_model_stats(model, input_sample['input_ids'].to(device))
